from data_loader import load_cifar10_data
from model import Cifar10ResNet
from train import train_model
from evaluate import evaluate_model
from visualize import visualize_results
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
import matplotlib.pyplot as plt

def plot_training_curves(train_losses, train_accuracies):
    """
    绘制训练曲线
    
    参数:
        train_losses (list): 训练损失列表
        train_accuracies (list): 训练准确率列表
    """
    plt.figure(figsize=(12, 5))
    
    # 绘制损失曲线
    plt.subplot(1, 2, 1)
    plt.plot(train_losses, label='Training Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Training Loss')
    plt.legend()
    
    # 绘制准确率曲线
    plt.subplot(1, 2, 2)
    plt.plot(train_accuracies, label='Training Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.title('Training Accuracy')
    plt.legend()
    
    plt.tight_layout()
    plt.savefig('training_curves.png')
    plt.show()

if __name__ == "__main__":
    # 加载数据
    train_loader, test_loader, class_names = load_cifar10_data(batch_size=128)
    
    # 初始化模型
    model = Cifar10ResNet()
    
    # 定义损失函数
    criterion = nn.CrossEntropyLoss()
    
    # 定义优化器 (使用AdamW，比Adam有更好的正则化效果)
    optimizer = optim.AdamW(
        model.parameters(), 
        lr=0.001, 
        weight_decay=0.0001,  # L2正则化
        betas=(0.9, 0.999)
    )
    
    # 定义学习率调度器 (余弦退火)
    scheduler = CosineAnnealingLR(
        optimizer, 
        T_max=50,  # 50个epoch
        eta_min=1e-6  # 最小学习率
    )
    
    # 训练模型
    print("开始训练...")
    trained_model, train_losses, train_accuracies = train_model(
        model, 
        train_loader, 
        criterion, 
        optimizer, 
        scheduler, 
        epochs=50
    )
    
    # 绘制训练曲线
    plot_training_curves(train_losses, train_accuracies)
    
    # 评估模型
    print("\n评估模型...")
    accuracy, report, cm = evaluate_model(
        trained_model, 
        test_loader, 
        class_names
    )
    
    # 可视化预测结果
    print("\n可视化预测结果...")
    visualize_results(
        trained_model, 
        test_loader, 
        class_names, 
        num_images=10
    )
    
    # 保存模型
    torch.save(trained_model.state_dict(), "cifar10_resnet.pth")
    print(f"\n模型已保存为 cifar10_resnet.pth，测试准确率: {accuracy:.2f}%")